package com.nbsaas.boot.thread;

import org.jetbrains.annotations.NotNull;
import org.springframework.core.task.TaskDecorator;

public class TenantAwareTaskDecorator implements TaskDecorator {
    @NotNull
    @Override
    public Runnable decorate(@NotNull Runnable runnable) {
        String tenant = TenantContext.getTenant();
        return () -> {
            TenantContext.setTenant(tenant);
            try {
                runnable.run();
            } finally {
                TenantContext.clear();
            }
        };
    }
}
